import qualified Data.Set as Set
import qualified Data.Map as Map
import Data.List ((\\))

type Jolt = Int
data Adapters = Adapters (Set.Set Jolt) Jolt deriving (Show)

parseInput :: String -> Adapters
parseInput string = 
    Adapters (Set.fromList $ (max+3):numbers) (max+3)
    where
        numbers = map read . lines $ string
        max = maximum numbers

type Record = Map.Map Jolt Integer

arrangementCount :: Adapters -> Integer
arrangementCount adapters@(Adapters numbers max) = 
    fst (backtrack initRecord max)
    where
        initRecord = Map.fromList [(1,1),(2,2),(3,4)]
        backtrack :: Record -> Jolt -> (Integer,Record)
        backtrack record jolt
            | not (Set.member jolt numbers) = (0, record)
            | otherwise                     = 
                case Map.lookup jolt record of
                    Just c -> (c,record)
                    Nothing -> (count,newRecord)
            where
                (count1, record') = backtrack record (jolt-1)
                (count2, record'') = backtrack record' (jolt-2)
                (count3, record''') = backtrack record'' (jolt-3)

                count = count1+count2+count3
                newRecord = Map.insert jolt count record'''

arrangementCount' :: Adapters -> Integer
arrangementCount' adapters@(Adapters numbers max) = 
    case Map.lookup max record of
        Just count -> count
    where
        numberLst = Set.toList numbers \\ [1,2]
        initRecord = Map.fromList [(1,1),(2,2)]

        record = foldl step initRecord numberLst

        getCount :: Record -> Jolt -> Integer
        getCount record jolt = 
            case Map.lookup jolt record of
                Just count -> count
                Nothing    -> 0

        step :: Record -> Jolt -> Record
        step record jolt = newRecord
            where
                count1 = getCount record (jolt-1)
                count2 = getCount record (jolt-2)
                count3 = getCount record (jolt-3)
                count = count1 + count2 + count3
                newRecord = Map.insert jolt count record

main :: IO ()
main = do
    input <- readFile "input.txt"
    print . arrangementCount . parseInput $ input
    print . arrangementCount' . parseInput $ input
